use std::{collections::HashMap, sync::Arc};

use async_trait::async_trait;
use log::{debug, info};
use russh::{
    server::{self, Auth, Handler, Msg, Server, Session},
    Channel, ChannelId, CryptoVec,
};
use tokio::sync::{broadcast, RwLock};
use yang2::{context::Context, data::DataTree};

use crate::{
    codec::{decode, encode},
    ctx::load_context,
    NcMsg, NetconfServerHandler, DEFAULT_SSH_PORT,
};

const DEFAULT_THREADS: u16 = 1;

// #[tokio::main]

pub async fn start_netconf_server(
    port: u16,
    threads: u16,
    ctx: Arc<Context>,
    module_handlers: Arc<HashMap<String, Arc<RwLock<Box<dyn ModuleHandler>>>>>,
    // notify_sender: tokio::sync::broadcast::Sender<NcMsg>,
) -> Result<(), anyhow::Error> {
    // let client_key = russh_keys::key::KeyPair::generate_ed25519().unwrap();
    // let client_pubkey = Arc::new(client_key.clone_public_key().unwrap());
    let mut config = russh::server::Config::default();
    config
        .keys
        .push(russh_keys::key::KeyPair::generate_ed25519().unwrap());
    let config = Arc::new(config);
    let (tx, rx) = broadcast::channel(100);
    let mut sh = SshServer {
        threads,
        idle_timeout: 0,
        endpoint: vec![],
        ctx,
        module_handlers,
        tx,
        rx,
    };
    sh.run_on_address(config, &format!("0.0.0.0:{}", port))
        .await?;
    Ok(())
}

pub async fn default_start_netconf_server(
    module_handlers: Arc<HashMap<String, Arc<RwLock<Box<dyn ModuleHandler>>>>>,
) -> Result<(), anyhow::Error> {
    let ctx = load_context()?;
    start_netconf_server(
        DEFAULT_SSH_PORT,
        DEFAULT_THREADS,
        ctx,
        module_handlers,
    )
    .await
}

pub struct SshServer {
    pub threads: u16,
    pub idle_timeout: u16,
    pub endpoint: Vec<NcEndpoint>,
    pub ctx: Arc<Context>,
    pub module_handlers: Arc<HashMap<String, Arc<RwLock<Box<dyn ModuleHandler>>>>>,
    rx: broadcast::Receiver<NcMsg>,
    tx: broadcast::Sender<NcMsg>,
}

#[async_trait]
#[allow(unused_variables)]
pub trait ModuleHandler: Send + 'static + Sync {
    /// module init method
    async fn init(&mut self, data_tree: &DataTree, notify_sender: broadcast::Sender<NcMsg>);

    /// module config method
    /// trigger by edit-config NETCONF operation
    async fn config(
        &mut self,
        data_tree: &DataTree,
        notify_sender: broadcast::Sender<NcMsg>,
    ) -> Result<(), anyhow::Error>;

    /// module rpc call method
    /// trigger by rpc NETCONF operation
    async fn rpc_call(
        &mut self,
        op_tree: Option<&DataTree>,
        notify_sender: broadcast::Sender<NcMsg>,
    ) -> Result<Option<DataTree>, anyhow::Error>;

    /// module get method
    /// trigger by get NETCONF operation
    async fn get(&self, data: &mut DataTree, notify_sender: broadcast::Sender<NcMsg>);

    /// module shutdown method
    /// server shutdown or shutdown by module manager
    async fn shutdown(&mut self, notify_sender: broadcast::Sender<NcMsg>) {
        info!("module shutdown...");
    }
}

#[derive(Clone, Default)]
pub struct NcEndpoint {
    pub name: String,
    pub ssh: NcEnpServerParams,
}
#[derive(Clone, Default)]
pub struct NcEnpServerParams {
    pub local_address: String,
    pub local_port: u16,
}

impl server::Server for SshServer {
    type Handler = NetconfHandler;
    fn new_client(&mut self, _: Option<std::net::SocketAddr>) -> NetconfHandler {
        NetconfHandler {
            channel_id: None,
            data_cache: vec![],
            threads: self.threads,
            ctx: self.ctx.clone(),
            tx: self.tx.clone(),
            rx: self.rx.resubscribe(),
        }
    }

    fn handle_session_error(&mut self, error: <Self::Handler as Handler>::Error) {
        debug!("error: {}", error.to_string());
    }
}

#[async_trait]
impl NetconfServerHandler for NetconfHandler {
    async fn on_message(&mut self, msg: crate::NcMsg) -> NcMsg {
        match msg {
            crate::NcMsg::Hello(_, d) => {
                info!("recv hello message: {:?}", d);
                return NcMsg::None;
            }
            _ => (),
        }
        NcMsg::None
    }

    async fn notify(&mut self, _msg: crate::NcMsg) {
        todo!()
    }

    async fn on_hello(&mut self, _msg: crate::NcMsg) -> NcMsg {
        todo!()
    }
}

pub struct NetconfHandler {
    pub channel_id: Option<ChannelId>,
    pub data_cache: Vec<u8>,
    pub threads: u16,
    pub ctx: Arc<Context>,
    pub tx: broadcast::Sender<NcMsg>,
    pub rx: broadcast::Receiver<NcMsg>,
}

/// NETCONF(RFC6242) message handler over SSH
/// The handler recieve NETCONF message, and then exchange the request/response message with datastore.
#[async_trait]
impl Handler for NetconfHandler {
    type Error = anyhow::Error;

    async fn channel_open_session(
        &mut self,
        channel: Channel<Msg>,
        _session: &mut Session,
    ) -> Result<bool, Self::Error> {
        let channel_id = channel.id();
        debug!("new session, channel_id: {}", channel_id);
        self.channel_id = Some(channel_id);
        // let (tx, mut rx) = broadcast::channel::<NcMsg>(100);
        let mut rx = self.rx.resubscribe();

        tokio::spawn(async move {
            if let Ok(nc_msg) = rx.recv().await {
                debug!("session handle thread recv a nc msg");
                let _ = channel.data(nc_msg.to_ncxml().as_bytes()).await;
            }
        });
        Ok(true)
    }

    async fn auth_password(&mut self, user: &str, password: &str) -> Result<Auth, Self::Error> {
        debug!("username: {}, password: {}", user, password);
        Ok(server::Auth::Accept)
    }

    async fn auth_succeeded(&mut self, _session: &mut Session) -> Result<(), Self::Error> {
        debug!("auth success");
        Ok(())
    }

    async fn subsystem_request(
        &mut self,
        channel: ChannelId,
        name: &str,
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        debug!("subsystem, name: {}", name);
        if "netconf".eq(name) {
            session.channel_success(channel);

            session.data(
                channel,
                CryptoVec::from_slice(
                    format!(
                        r#"
            <?xml version="1.0" encoding="UTF-8"?>
            <hello>
              <capabilities>
                <capability>
                  urn:ietf:params:xml:ns:netconf:base:1.0
                </capability>
                <capability>
                  urn:ietf:params:ns:netconf:capability:startup:1.0
                </capability>
              </capabilities>
              <session-id>{}</session-id>
            </hello>
            ]]>]]>
            "#,
                        channel.to_string()
                    )
                    .as_bytes(),
                ),
            );
            return Ok(());
        } else {
            session.channel_failure(channel);
            return Ok(());
        }
    }

    async fn channel_eof(
        &mut self,
        _channel: ChannelId,
        _session: &mut Session,
    ) -> Result<(), Self::Error> {
        debug!("channel eof");
        self.data_cache.clear();
        Ok(())
    }

    async fn data(
        &mut self,
        channel: ChannelId,
        data: &[u8],
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        debug!(
            "data comming, channel_id: {}, data_len: {}",
            channel,
            data.len()
        );
        self.data_cache.append(&mut data.to_vec());
        let (pos, msgs) = decode(&self.data_cache);
        let (_, r) = self.data_cache.split_at(pos);
        self.data_cache = r.to_vec();
        {
            if let Some(msgs) = &msgs {
                for msg in msgs {
                    let re = self.on_message(msg.clone()).await;
                    match re {
                        NcMsg::None => (),
                        default => session.data(channel, CryptoVec::from_slice(&encode(default))),
                    }
                }
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::{collections::HashMap, sync::Arc};

    use crate::ssh::default_start_netconf_server;
    #[tokio::test]
    async fn test_run() -> Result<(), anyhow::Error> {
        env_logger::init();
        // let (tx, _rx) = tokio::sync::broadcast::channel(10);
        default_start_netconf_server(Arc::new(HashMap::new())).await
    }
}
